#include "crfHolder0.h"
#define INVALID_TAG -1


void crfHolder0::build_lattice(sentence& sen)
{
	_map_pool.free();
	_link_pool.free();
	_tagid2gridid.resize(sen.l);
	int total_grid_num=0;
	for(int i=0;i<sen.l;i++) {
		map_int& mpi=_tagid2gridid[i];
		mpi.init(_map_pool,sen.word_feas_num[i]*2);
		total_grid_num+=sen.word_feas_num[i];
	}
	_grid_memPool.alloc(total_grid_num);
	grid_node0* grid_node0_buf = (grid_node0*)_grid_memPool;
	_node_grid.resize(sen.l);
	for(int i=0;i<total_grid_num;i++)
		grid_node0_buf[i].init();

	for(int i=0;i<sen.l;i++) {
		_node_grid[i]=grid_node0_buf;
		grid_node0_buf+=sen.word_feas_num[i];
	}
		
	for(int i=0;i<sen.l;i++) {
		map_int* tagid_cur      = &_tagid2gridid[i];
		grid_node0* gridnode_cur= _node_grid[i];

		int icount = 0;
		for(int ifs=0;ifs<sen.word_feas_num[i];ifs++)
		{
			span_feature* spn= &sen.word_feas[i][ifs];
			int idx_gd = (*tagid_cur)[spn->y_cur];
			if(idx_gd==INVALID_TAG)
			{
				grid_node0& gd2= gridnode_cur[icount];
				tagid_cur->insert(spn->y_cur,icount);
				gd2.tag_id = spn->y_cur;
				gd2.push_feas(_link_pool,spn);
				icount++;
			}
			else
			{
				grid_node0& gd2= gridnode_cur[idx_gd];
				assert_loc( gd2.tag_id == spn->y_cur);
				gd2.push_feas(_link_pool,spn);
			}
		}
	}

	_double_pool.alloc(sen.l*2+4);
	_alpha = (double*)_double_pool;
	_beta = _alpha+sen.l+2;
};

void crfHolder0::calculate_alpha(sentence& sen)
{
	for(int i=0;i<sen.l;i++) {
		map_int &cur_mapid = _tagid2gridid[i];
		if(cur_mapid.size()==0) {
			_alpha[i]=0;
			continue;
		}

		double cur_p = -1;
		//-------------
		for(int it=0;it<_tag_size;it++)
		{
			int layerid=cur_mapid[it];
			if(layerid==INVALID_TAG)	continue;
			grid_node0* gdnd = &_node_grid[i][layerid];

			double p = 0;
			zero_feature_link* inlink = gdnd->node_feas;
			while(inlink) {
				span_feature* spn=inlink->cur_span;
				p += _lambda[spn->index];
				inlink = inlink->next;
			}
			gdnd->grid_alpha = p;

			if(cur_p<0)
				cur_p = p;
			else
				cur_p= log_sum_exp( cur_p, p );
		}
		//-------------
		if(cur_p>=0)
			_alpha[i]= cur_p;
		else
			_alpha[i]= 0;
	}
	_allZ=0;
	for(int i=0;i<sen.l;i++) {
		_allZ+=_alpha[i];
	}

}


double crfHolder0::calculate_cost(sentence& sen)
{
	double p_all=0;
	for(int i=0;i<sen.l;i++)
	{
		int it= sen.labels[i];
		int layerid=_tagid2gridid[i][it];
		if(layerid==INVALID_TAG)	continue;
		grid_node0* gdnd = &_node_grid[i][layerid];
		double p = gdnd->grid_alpha;

		if(p>=0)
			p_all += p;
	}
	return p_all;
}
 

void crfHolder0::calculate_gradient(sentence& sen,double& z_all)
{
	z_all=_allZ;
	assert_loc(z_all!=-1);

	for(int i=0;i<sen.l;i++){
		map_int &cur_mapid = _tagid2gridid[i];
		grid_node0* cur_grids = _node_grid[i];
		//-------------
		for(int it=0;it<_tag_size;it++)
		{
			int layerid=cur_mapid[it];
			if(layerid==INVALID_TAG)	continue;
			grid_node0* gdnd = &cur_grids[layerid];
			zero_feature_link* inlink = gdnd->node_feas;
			double post_it = exp(gdnd->grid_alpha - _alpha[i]) ;
			while(inlink) {
				span_feature* inedg=inlink->cur_span;
				_gradient[inedg->index] += post_it;
				inlink=inlink->next;
			}
		}
		//---------------------
	}

	for(int i=0;i<sen.l;i++)
	{
		int it = sen.labels[i];
		for(int ifea=0;ifea<sen.word_feas_num[i];ifea++)
		{
			span_feature& spn = sen.word_feas[i][ifea];
			if(spn.y_cur == it )
			{
				_gradient[spn.index]--;
			}
		}
	}
}

double crfHolder0::process_sen(sentence& sen)
{
	/*for(int i=0;i<sen.l;i++)
	{
		for(int il=0;il<sen.word_feas_num[i];il++)
		{
			span_feature& spn = sen.word_feas[i][il];
			spn.y_cur++;
			spn.y_pre++;
		}
		sen.labels[i]++;
	}*/

	double all_z = -1;
	build_lattice(sen);
	calculate_alpha(sen);
	calculate_gradient(sen,all_z);
	double cost = calculate_cost(sen);
	assert_loc(all_z>=0);
	return all_z - cost;
}

double crfHolder0::process_decode(sentence& sen,vector< vector<int> >& predict_label,vector< vector<double> >& predict_prob)
{
	/*for(int i=0;i<sen.l;i++)
	{
		for(int il=0;il<sen.word_feas_num[i];il++)
		{
			span_feature& spn = sen.word_feas[i][il];
			spn.y_cur++;
			spn.y_pre++;
		}
		sen.labels[i]++;
	}*/
	build_lattice(sen);
	calculate_alpha(sen);
	return decode(sen,predict_label,predict_prob);
}


struct decode_pair
{
	int first;
	double second;
};

bool inverse_cmp_my(const decode_pair& s, const decode_pair& t)
{
   return s.second>t.second;
}

double crfHolder0::decode(sentence& sen,vector< vector<int> >& predict_label,vector< vector<double> >& predict_prob)
{
	const int nbest = 10;
	predict_label.clear();
	predict_label.resize(sen.l);
	predict_prob.clear();
	predict_prob.resize(sen.l);
	int  last_it = -1;
	double total_p = 0;
	_decode_memPool.alloc(sen.l*nbest*sizeof(decode_pair));
	decode_pair* intmemPool = (decode_pair*)(char*)_decode_memPool;
	memset( intmemPool, 0, sen.l*nbest*sizeof(decode_pair) );
	pair<decode_pair*,decode_pair*> ip;

	for(int i=0;i<sen.l;i++){
		predict_label[i].resize(nbest,-1);
		predict_prob[i].resize(nbest,-1);
		map_int &cur_mapid = _tagid2gridid[i];
		grid_node0* cur_grids = _node_grid[i];
		decode_pair*  s_pool = intmemPool+i*nbest;
		if(cur_mapid.size()==0)
		{
			for(int ii = nbest-1;ii>=0;ii--){
				decode_pair* tpdest= s_pool+ii;
				tpdest->first = -1;
				tpdest->second = 0;
			}
			continue;
		}
		//-------------
		last_it = -1;
		double max_p = -1;
		decode_pair tmp_v;
		for(int it=0;it<_tag_size;it++)
		{
			int layerid=cur_mapid[it];
			if(layerid==INVALID_TAG)	continue;
			grid_node0* gdnd = &cur_grids[layerid];
			tmp_v.second = gdnd->grid_alpha;tmp_v.first = it;
			ip = equal_range(s_pool, s_pool+nbest, tmp_v, inverse_cmp_my);
			int pos=ip.second-s_pool;
			if(pos>=nbest) continue;

			for(int ii = nbest-1;ii>pos;ii--){
				decode_pair* tpdest= s_pool+ii;
				decode_pair* tpsrc= s_pool+ii-1;
				tpdest->first = tpsrc->first;
				tpdest->second = tpsrc->second;
			}
			decode_pair* tpdest= s_pool+pos;
			tpdest->first = it;
			tpdest->second = gdnd->grid_alpha;
		}
	}
	for(int ib=0;ib<nbest;ib++){
		if(ib==0)
			for(int i=0;i<sen.l;i++){
				predict_label[i][ib] = intmemPool[i*nbest].first;
				predict_prob[i][ib] = exp(intmemPool[i*nbest].second-_alpha[i]);
				total_p+= intmemPool[i*nbest].second;
			}
		else
			for(int i=0;i<sen.l;i++){
				predict_label[i][ib] = intmemPool[i*nbest+ib].first;
				predict_prob[i][ib] = exp(intmemPool[i*nbest+ib].second-_alpha[i]);
			}
	}
	return total_p;

	for(int ib=0;ib<nbest;ib++){
		if(ib==0)
			for(int i=0;i<sen.l;i++){
				predict_label[i][ib] = intmemPool[i*nbest].first;
				total_p+= intmemPool[i*nbest].second;
			}
		else
		{
			double min_p = DBL_MAX;
			double min_it = -1;
			for(int i=0;i<sen.l;i++){
				int curi=predict_label[i][ib-1];
				if(curi>=nbest-1)
					continue;
				double dv = intmemPool[i*nbest+curi].second-intmemPool[i*nbest+curi+1].second;
				assert_loc(dv>=0);
				if(dv>min_p) {
					min_p = dv;
					min_it = i;
				}
			}
			for(int i=0;i<sen.l;i++){
				if(i!=min_it)
					predict_label[i][ib] = predict_label[i][ib-1];
				else {
					int curi=predict_label[i][ib-1];
					predict_label[i][ib] = intmemPool[i*nbest+curi+1].first;
				}
			}
		}
	}

	return total_p;
}

